import triton
import torch
import torch_npu
import triton.language as tl

from torch._inductor.runtime import triton_helpers
from torch._inductor.runtime.triton_helpers import libdevice

from .util import triton_helper_compile, compiled_kernel_list

@triton.jit
def fused_kernel1(in_ptr0, in_ptr1, in_ptr2, out_ptr1, xnumel, rnumel):
    xnumel = 9
    XBLOCK: tl.constexpr = 1
    rnumel = 2048
    RBLOCK: tl.constexpr = 2048
    xoffset = tl.program_id(0) * XBLOCK
    xindex = tl.full([1], xoffset, tl.int32)
    xmask = tl.full([RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[:]
    roffset = 0
    rmask = tl.full([RBLOCK], True, tl.int1)
    r1 = rindex
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (r1 + (2048*x0)), None)
    tmp1 = tl.load(in_ptr1 + (r1 + (2048*x0)), None)
    tmp7 = tl.load(in_ptr2 + (r1), None, eviction_policy='evict_last')
    tmp2 = tmp0 + tmp1
    tmp3 = tmp2 * tmp2
    tmp4 = tl.broadcast_to(tmp3, [RBLOCK])
    tmp6 = triton_helpers.promote_to_tensor(tl.sum(tmp4, 0))
    tmp8 = 2048.0
    tmp9 = tmp6 / tmp8
    tmp10 = 1e-05
    tmp11 = tmp9 + tmp10
    tmp12 = libdevice.rsqrt(tmp11)
    tmp13 = tmp2 * tmp12
    tmp14 = tmp7 * tmp13
    tl.store(out_ptr1 + (r1 + (2048*x0)), tmp14, None)

def test_compile_for_fused_kernel1():
    signature = {0: '*fp32', 1: '*fp32', 2: '*fp32', 3: '*fp32', 4: 'i32', 5: 'i32'}
    triton_helper_compile(func=fused_kernel1, signature=signature)

compiled_kernel_list["fused_kernel1"] = test_compile_for_fused_kernel1